#include "znx_session_pool.h"
#include "comm/znx_crc.h"

static znx_bool_t znx_t_ession_pool_init_handle(void *data);
static znx_bool_t znx_t_session_pool_proc_exit_handle(void *data);
static znx_bool_t znx_t_session_pool_quit_handle(void *data);

size_t g_session_pool_max_conns = 0;
__thread znx_session_pool_t *t_session_pool = NULL;


void
znx_t_session_pool_init(size_t max_conns)
{
    if (max_conns == 0) {
        return;
    }

    g_session_pool_max_conns = max_conns;
    znx_worker_processor_add_init_hook(znx_t_ession_pool_init_handle);
    znx_worker_processor_add_proc_exit_hook(znx_t_session_pool_proc_exit_handle);
    znx_worker_processor_add_quit_hook(znx_t_session_pool_quit_handle);
}


znx_session_pool_t *
znx_session_pool_create(size_t max_conns)
{
    if (max_conns == 0) {
        return NULL;
    }

    znx_session_pool_t *pool = znx_thread_calloc(sizeof(znx_session_pool_t));
    pool->max_conns = max_conns;
    znx_queue_init(&pool->lru_head);
    znx_rbtree_init(&pool->tree, &pool->tree_sentinel, znx_session_item_cmp);

    return pool;
}


static void
znx_session_pool_delete_item(znx_session_pool_t *pool,
    znx_session_item_t *item)
{
    znx_queue_remove(&item->lru_queue);
    znx_rbtree_delete(&pool->tree, &item->tree_node);
    pool->cur_conns--;

    item->free_handle(item->conn);
    znx_thread_free(item);
}


static void
znx_session_conn_read_handle(znx_event_t *event)
{
    znx_conn_t *conn = event->data;
    znx_session_item_t *item = conn->data;
    znx_session_pool_t *pool = item->pool;
    znx_session_pool_delete_item(pool, item);
}


void
znx_session_pool_put(znx_session_pool_t *pool,
    znx_str_t *tag, znx_conn_t *conn, int64_t timeout_ms,
    znx_conn_handle_pt free_handle)
{
    znx_address_t       *peer_address;
    znx_session_item_t  *item;
    znx_queue_t         *q;
    size_t              peer_address_len = 0;
    size_t              item_len = 0;

    if (free_handle == NULL) {
        free_handle = znx_session_pool_default_free_conn;
    }

    if (timeout_ms <= 0) {
        free_handle(conn);
        return;
    }

    if (conn->fd == -1 ||
        conn->read.eof || conn->read.error ||
        conn->write.eof || conn->write.error)
    {
        free_handle(conn);
        return;
    }

    peer_address = znx_conn_peer_address(conn);
    if (proc_exit || peer_address == NULL) {
        free_handle(conn);
        return;
    }

    if (peer_address->sockaddr.sa_family != AF_INET &&
        peer_address->sockaddr.sa_family != AF_INET6 &&
        peer_address->sockaddr.sa_family != AF_UNIX)
    {
        free_handle(conn);
        return;
    }

    // some readers put conn into the session pool after reading
    // the business data but not reading EAGAIN.
    // at this point the ready flag is not cleared.
    // we will receive read event immediately when
    // we execute znx_event_loop_enable_read.
    if (conn->read.ready != 0) {
        char buf[4];
        znx_result_t    ret;
        ret = znx_conn_read(conn, buf, 4);
        if (ret != ZNX_AGAIN) {
            free_handle(conn);
            return;
        }
    }

    if (!znx_event_loop_start_io(conn)) {
        free_handle(conn);
        return;
    }

    conn->read.handler = znx_session_conn_read_handle;
    if (!znx_event_loop_enable_read(conn)) {
        free_handle(conn);
        return;
    }

    znx_event_loop_add_timer(&conn->read, timeout_ms);

    if (pool->cur_conns >= pool->max_conns) {
        q = znx_queue_last(&pool->lru_head);
        item = znx_queue_data(q, znx_session_item_t, lru_queue);
        znx_session_pool_delete_item(pool, item);
    }

    peer_address_len = (size_t)znx_address_socklen(peer_address);

    // build item key.
    // item + tag + address
    item_len = sizeof(znx_session_item_t) + tag->len + peer_address_len;
    item = znx_thread_calloc(item_len);
    item->key.data = (u_char *)item + sizeof(znx_session_item_t);
    znx_memcpy(item->key.data, tag->data, tag->len);
    znx_memcpy(item->key.data + tag->len, (u_char *)peer_address, peer_address_len);
    item->key.len = tag->len + peer_address_len;
    item->key_crc = znx_crc32_long(item->key.data, item->key.len);

    item->conn = conn;
    item->pool = pool;
    item->free_handle = free_handle;
    conn->data = item;

    znx_queue_insert_head(&pool->lru_head, &item->lru_queue);
    znx_rbtree_insert(&pool->tree, &item->tree_node);
    pool->cur_conns++;
}


znx_bool_t
znx_session_conn_avalive_check(znx_conn_t *conn)
{
    int         err = 0;
    socklen_t   len = 0;

    if (getsockopt(conn->fd, SOL_SOCKET, SO_ERROR, (void *)&err, &len) == -1) {
        return ZNX_FALSE;
    }

    if (err) {
        conn->read.error = 1;
        conn->write.error = 1;
        return ZNX_FALSE;
    }

    return ZNX_TRUE;
}


// the address must be reinitialized after bero.
static znx_conn_t *
znx_session_pool_get_helper(znx_session_pool_t *pool,
    znx_str_t *tag, znx_address_t *address)
{
    size_t              key_len;
    size_t              address_len;
    znx_session_item_t  *item;
    znx_rbtree_node_t   *node;
    znx_conn_t          *conn;

    address_len = (size_t)znx_address_socklen(address);

    key_len = tag->len + address_len;
    znx_t_buf_grow(key_len);
    znx_memcpy(t_buf, tag->data, tag->len);
    znx_memcpy(t_buf + tag->len, (u_char *)address, address_len);

    znx_session_item_t  target;
    target.key.data = t_buf;
    target.key.len = key_len;
    target.key_crc = znx_crc32_long(t_buf, key_len);

    node = znx_rbtree_lookup(&pool->tree, &target.tree_node);
    if (node == NULL) {
        return NULL;
    }

    item = (znx_session_item_t *)node;

    znx_rbtree_delete(&pool->tree, &item->tree_node);
    znx_queue_remove(&item->lru_queue);
    pool->cur_conns--;

    conn = item->conn;
    conn->data = (void *)item->free_handle;
    conn->read.handler = NULL;
    znx_event_loop_remove_conn(conn);

    znx_thread_free(item);

    return conn;
}


znx_conn_t *
znx_session_pool_get(znx_session_pool_t *pool,
    znx_str_t *tag, znx_address_t *address)
{
    znx_conn_t                  *conn;
    znx_conn_handle_pt          free_handle;

    while (1) {
        conn = znx_session_pool_get_helper(pool, tag, address);
        if (conn == NULL) {
            return NULL;
        }

        if (!znx_session_conn_avalive_check(conn)) {
            free_handle = (znx_conn_handle_pt)conn->data;
            free_handle(conn);
            continue;
        }

        conn->data = NULL;
        return conn;
    }

    // unreachable
}


void
znx_session_pool_clean(znx_session_pool_t *pool)
{
    znx_queue_t         *q;
    znx_session_item_t  *item;

    while (pool->cur_conns > 0) {
        q = znx_queue_last(&pool->lru_head);
        item = znx_queue_data(q, znx_session_item_t, lru_queue);
        znx_session_pool_delete_item(pool, item);
    }
}


void
znx_session_pool_destroy(znx_session_pool_t *pool)
{
    znx_session_pool_clean(pool);
    znx_thread_free(pool);
}


static znx_bool_t
znx_t_ession_pool_init_handle(void *data)
{
    t_session_pool = znx_session_pool_create(g_session_pool_max_conns);
    return ZNX_TRUE;
}


static znx_bool_t
znx_t_session_pool_proc_exit_handle(void *data)
{
    if (t_session_pool == NULL) {
        return ZNX_TRUE;
    }

    znx_session_pool_clean(t_session_pool);
    return ZNX_TRUE;
}


static znx_bool_t
znx_t_session_pool_quit_handle(void *data)
{
    if (t_session_pool == NULL) {
        return ZNX_TRUE;
    }

    znx_session_pool_destroy(t_session_pool);
    return ZNX_TRUE;
}
